/*
 *    Copyright 2011-2012 The Hook Team
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 */
package org.hook.core.mongo;

import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.mongodb.core.query.Update;

import com.mongodb.DB;
import com.mongodb.DBAddress;
import com.mongodb.DBCollection;
import com.mongodb.Mongo;
import com.mongodb.MongoOptions;

/**
 * 集成SPRING-MONGO模板(扩展支持多mongodb数据源)
 * 
 * @author suwuqiang 2011-8-22
 */
public class MongoDynamicDAO implements MongoDAO, InitializingBean {
	private static Logger logger = LoggerFactory.getLogger(MongoDynamicDAO.class);
	private static final String MONGO_URL_PREFIX = "mongo://";
	private Map<String, MongoDataSource> datasources;
	private MongoDataSource dataSource;
	protected MongoTemplate mongoTemplate;

	public MongoTemplate getMongoTemplate() {
		if (null == mongoTemplate || null == dataSource) {
			// service层显式调用MongoDataSourceContextHolder.setMongoDataSourceType方法，决定当前线程所需要的mongoDataSource
			String dataSourceKey = MongoDataSourceContextHolder.getMongoDataSourceType();
			MongoDataSource mongoDataSource = this.datasources.get(dataSourceKey);
			if (null == mongoDataSource) {
				throw new IllegalArgumentException("Target 'datasource:" + dataSourceKey + "' is not exists!");
			}
			this.setDataSource(mongoDataSource);
			// parse MongoDataSource
			String url = mongoDataSource.getUrl();
			String database = url.substring(url.lastIndexOf("/") + 1);
			try {
				MongoOptions options = new MongoOptions();
				DBAddress address = new DBAddress(url.substring(MONGO_URL_PREFIX.length()));
				Mongo mongo = new Mongo(address, options);
				mongoTemplate = new MongoTemplate(mongo, database);
			} catch (Exception e) {
				logger.error("Create MongoTemplate failure!", e);
				throw new RuntimeException("Create MongoTemplate failure!", e);
			}
		}
		return mongoTemplate;
	}

	public void setMongoTemplate(MongoTemplate mongoTemplate) {
		this.mongoTemplate = mongoTemplate;
	}

	@Override
	public void afterPropertiesSet() throws Exception {
		if (datasources == null || datasources.size() < 1) {
			throw new IllegalArgumentException("Property 'datasources' is required.");
		}
	}

	@Override
	public String getCollectionName(Class<?> entityClass) {
		return getMongoTemplate().getCollectionName(entityClass);
	}

	@Override
	public DBCollection getCollection(Class<?> entityClass) {
		String collectionName = this.getCollectionName(entityClass);
		return this.getCollection(collectionName);
	}

	@Override
	public DBCollection getCollection(String collectionName) {
		DBCollection dbCollection = null;
		if (null == collectionName || collectionName.equals("")) {
			throw new IllegalArgumentException("Property 'collectionName:" + collectionName + "' is illegal.");
		}
		if (getMongoTemplate().collectionExists(collectionName)) {
			dbCollection = getMongoTemplate().getCollection(collectionName);
		}
		return dbCollection;
	}

	@Override
	public DB getDb() {
		return getMongoTemplate().getDb();
	}

	@Override
	public void dropCollection(String collectionName) {
		getMongoTemplate().dropCollection(collectionName);
	}

	@Override
	public void dropCollection(Class<?> entityClass) {
		getMongoTemplate().dropCollection(entityClass);
	}

	@Override
	public Set<String> getCollectionNames() {
		return getMongoTemplate().getCollectionNames();
	}

	@Override
	public long getCollections() {
		return getCollectionNames().size();
	}

	@Override
	public void insert(Object objectToSave) {
		getMongoTemplate().insert(objectToSave);
	}

	@Override
	public void insert(Collection<? extends Object> batchToSave, Class<?> entityClass) {
		getMongoTemplate().insert(batchToSave, entityClass);
	}

	@Override
	public void insert(Object objectToSave, String collectionName) {
		getMongoTemplate().insert(objectToSave, collectionName);
	}

	@Override
	public void insert(Collection<? extends Object> batchToSave, String collectionName) {
		getMongoTemplate().insert(batchToSave, collectionName);
	}

	@Override
	public void remove(Object obj) {
		getMongoTemplate().remove(obj);
	}

	@Override
	public void remove(final Query query, String collectionName) {
		getMongoTemplate().remove(query, collectionName);
	}

	@Override
	public void remove(final Query query, Class<?> entityClass) {
		getMongoTemplate().remove(query, entityClass);
	}

	@Override
	public void update(Query query, Update update, Class<?> entityClass) {
		getMongoTemplate().updateFirst(query, update, entityClass);
	}

	@Override
	public void update(Query query, Update update, String collectionName) {
		getMongoTemplate().updateFirst(query, update, collectionName);
	}

	@Override
	public void updateMulti(Query query, Update update, Class<?> entityClass) {
		getMongoTemplate().updateMulti(query, update, entityClass);
	}

	@Override
	public void updateMulti(Query query, Update update, String collectionName) {
		getMongoTemplate().updateMulti(query, update, collectionName);
	}

	@Override
	public <T> T load(Object id, Class<T> entityClass) {
		return getMongoTemplate().findById(id, entityClass);
	}

	@Override
	public <T> T load(Object id, Class<T> entityClass, String collectionName) {
		return getMongoTemplate().findById(id, entityClass, collectionName);
	}

	@Override
	public <T> List<T> find(Class<T> entityClass) {
		return getMongoTemplate().findAll(entityClass);
	}

	@Override
	public <T> List<T> find(Class<T> entityClass, String collectionName) {
		return getMongoTemplate().findAll(entityClass, collectionName);
	}

	@Override
	public <T> List<T> find(final Query query, Class<T> entityClass) {
		return getMongoTemplate().find(query, entityClass);
	}

	@Override
	public <T> List<T> find(final Query query, Class<T> entityClass, String collectionName) {
		return getMongoTemplate().find(query, entityClass, collectionName);
	}

	@Override
	public <T> T findOne(final Query query, Class<T> entityClass) {
		return getMongoTemplate().findOne(query, entityClass);
	}

	@Override
	public <T> T findOne(final Query query, Class<T> entityClass, String collectionName) {
		return getMongoTemplate().findOne(query, entityClass, collectionName);
	}

	@Override
	public long count(String collectionName) {
		return getCollection(collectionName).count();
	}

	@Override
	public long count(final Query query, String collectionName) {
		return getCollection(collectionName).getCount(query.getQueryObject(), query.getFieldsObject());
	}

	@Override
	public long count(final Query query, String collectionName, boolean limit) {
		if (!limit)
			return 0;
		return getCollection(collectionName).getCount(query.getQueryObject(), query.getFieldsObject(), query.getLimit(), query.getSkip());
	}

	public void setDataSource(MongoDataSource dataSource) {
		this.dataSource = dataSource;
	}

	public void setDatasources(Map<String, MongoDataSource> datasources) {
		this.datasources = datasources;
	}
}
